#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;                                     // tensor num for each queue

template<typename typeT>
class KernelAsinhGrad {
public:
    __aicore__ inline KernelAsinhGrad() {}
    __aicore__ inline void Init(GM_ADDR y, GM_ADDR dy, GM_ADDR z, uint32_t inputNum, 
                                uint32_t tileLength) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->inputNum = inputNum;
        this->tileLength = tileLength;

        yGm.SetGlobalBuffer((__gm__ typeT*)y, this->inputNum + 16);
        dyGm.SetGlobalBuffer((__gm__ typeT*)dy, this->inputNum + 16);
        zGm.SetGlobalBuffer((__gm__ typeT*)z, this->inputNum + 16);

        // pipe alloc memory to queue, the unit is Bytes
        pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileLength * sizeof(typeT));
        pipe.InitBuffer(inQueueDy, BUFFER_NUM, this->tileLength * sizeof(typeT));
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(typeT));

        pipe.InitBuffer(tmp1, this->tileLength * sizeof(float));
        pipe.InitBuffer(tmp2, this->tileLength * sizeof(float));
        
        //pipe.InitBuffer(tmp3, this->tileLength * sizeof(int8_t) / 8);
    }
    __aicore__ inline void Process() {
        int32_t loopCount = this->inputNum / this->tileLength;
        if(loopCount == 0) {
            uint32_t length = (this->inputNum + 15) / 16 * 16;
            CopyIn(0, length);
            Compute(0, this->inputNum);
            CopyOut(0, length);
        } else {
            for (int32_t i = 0; i < loopCount; i++) {
                CopyIn(i, this->tileLength);
                Compute(i, this->tileLength);
                CopyOut(i, this->tileLength);
            } 
            if (this->inputNum % this->tileLength > 0) {
                uint32_t length = this->inputNum % this->tileLength;
                uint32_t lengthalign = (length + 15) / 16 * 16;
                CopyIn(loopCount, lengthalign);
                Compute(loopCount, length);
                CopyOut(loopCount, lengthalign);
            }   
        }        
    }

private:
    __aicore__ inline void CopyIn(int32_t progress, uint32_t length) {
        LocalTensor<typeT> yLocal = inQueueY.AllocTensor<typeT>();
        DataCopy(yLocal, yGm[progress * this->tileLength], length);
        inQueueY.EnQue(yLocal);
        LocalTensor<typeT> dyLocal = inQueueDy.AllocTensor<typeT>();
        DataCopy(dyLocal, dyGm[progress * this->tileLength], length);
        inQueueDy.EnQue(dyLocal);
    }
    __aicore__ inline void Compute(int32_t progress, uint32_t length) {
        LocalTensor<typeT> yLocal = inQueueY.DeQue<typeT>();
        LocalTensor<typeT> dyLocal = inQueueDy.DeQue<typeT>();
        LocalTensor<typeT> zLocal = outQueueZ.AllocTensor<typeT>();

        if constexpr (std::is_same_v<typeT, float>) {

            // Exp(yLocal, yLocal, length);
            // Mul(dyLocal, dyLocal, yLocal, length);
            // Muls(dyLocal, dyLocal, typeT(2), length);
            // Mul(yLocal, yLocal, yLocal, length);
            // Adds(yLocal, yLocal, typeT(1), length);
            // Div(zLocal, dyLocal, yLocal, length);

            // Duplicate(zLocal, typeT(2), length);
            // Exp(yLocal, yLocal, length);
            // Mul(dyLocal, dyLocal, yLocal, length);            
            // Mul(dyLocal, dyLocal, zLocal, length);
            // Mul(yLocal, yLocal, yLocal, length);
            // Duplicate(zLocal, typeT(1), length);
            // Add(yLocal, yLocal, zLocal, length);
            // Div(zLocal, dyLocal, yLocal, length);

            Exp(yLocal, yLocal, length);
            Mul(dyLocal, dyLocal, yLocal, length); 
            Mul(yLocal, yLocal, yLocal, length);
            Muls(dyLocal, dyLocal, typeT(2), length);
            Adds(yLocal, yLocal, typeT(1), length);
            Div(zLocal, dyLocal, yLocal, length);

 
        } else if constexpr (std::is_same_v<typeT, half>) {
            auto yLocalf32 = tmp1.Get<float>();
            auto dyLocalf32 = tmp2.Get<float>();
            Cast(yLocalf32, yLocal, RoundMode::CAST_NONE, length);
            Cast(dyLocalf32, dyLocal, RoundMode::CAST_NONE, length);
            Exp(yLocalf32, yLocalf32, length);
            Mul(dyLocalf32, dyLocalf32, yLocalf32, length);
            Mul(yLocalf32, yLocalf32, yLocalf32, length);
            Muls(dyLocalf32, dyLocalf32, float(2), length);
            Adds(yLocalf32, yLocalf32, float(1), length);
            Div(dyLocalf32, dyLocalf32, yLocalf32, length);
            Cast(zLocal, dyLocalf32, RoundMode::CAST_ROUND, length);

            // Duplicate(zLocal, typeT(2), length);
            // Exp(yLocal, yLocal, length);
            // Mul(dyLocal, dyLocal, yLocal, length);            
            // Mul(dyLocal, dyLocal, zLocal, length);
            // Mul(yLocal, yLocal, yLocal, length);
            // Duplicate(zLocal, typeT(1), length);
            // Add(yLocal, yLocal, zLocal, length);
            // Div(zLocal, dyLocal, yLocal, length);

            // auto yLocalf32 = tmp1.Get<float>();
            // Cast(yLocalf32, yLocal, RoundMode::CAST_NONE, length);
            // Exp(yLocalf32, yLocalf32, length);
            // Cast(yLocal, yLocalf32, RoundMode::CAST_TRUNC, length);   //exp(y)
            // Mul(dyLocal, dyLocal, yLocal, length);     // dy = dy * exp(y)
            // Mul(yLocalf32, yLocalf32, yLocalf32, length);  
            // Cast(yLocal, yLocalf32, RoundMode::CAST_TRUNC, length);   // exp(y)^2  

            // Muls(dyLocal, dyLocal, typeT(2), length);
            // Adds(yLocal, yLocal, typeT(1), length);
            // Div(zLocal, dyLocal, yLocal, length);

            // auto yLocalf32 = tmp1.Get<float>();
            // Exp(yLocal, yLocal, length);
            // Mul(dyLocal, dyLocal, yLocal, length); 
            // Cast(yLocalf32, yLocal, RoundMode::CAST_NONE, length);
            // Mul(yLocalf32, yLocalf32, yLocalf32, length);
            // Cast(yLocal, yLocalf32, RoundMode::CAST_ROUND, length);
            // Muls(dyLocal, dyLocal, typeT(2), length);
            // Adds(yLocal, yLocal, typeT(1), length);
            // Div(zLocal, dyLocal, yLocal, length);   

            // auto YLocalf32 = tmp1.Get<float>();
            // auto DyLocalf32 = tmp2.Get<float>();             
            // Exp(yLocal, yLocal, length);
            // Mul(dyLocal, dyLocal, yLocal, length); 
            // Mul(yLocal, yLocal, yLocal, length);
            // Muls(dyLocal, dyLocal, typeT(2), length);
            // Adds(yLocal, yLocal, typeT(1), length);
            // Cast(YLocalf32, yLocal, RoundMode::CAST_NONE, length);
            // Cast(DyLocalf32, dyLocal, RoundMode::CAST_NONE, length);
            // Div(YLocalf32, DyLocalf32, YLocalf32, length);      
            // Cast(zLocal, YLocalf32, RoundMode::CAST_ROUND, length);
        }

        outQueueZ.EnQue<typeT>(zLocal);
        inQueueY.FreeTensor(yLocal);
        inQueueDy.FreeTensor(dyLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress, uint32_t length) {
        LocalTensor<typeT> zLocal = outQueueZ.DeQue<typeT>();
        DataCopy(zGm[progress * this->tileLength], zLocal, length);
        outQueueZ.FreeTensor(zLocal);
    }

private:
    TPipe pipe;
    TBuf<QuePosition::VECCALC> tmp1, tmp2;
    //TBuf<QuePosition::VECCALC> tmp1;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueY;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueDy;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueZ;
    GlobalTensor<typeT> yGm;
    GlobalTensor<typeT> dyGm;
    GlobalTensor<typeT> zGm;
    uint32_t inputNum;
    uint32_t tileLength;

};

extern "C" __global__ __aicore__ void asinh_grad(GM_ADDR y, GM_ADDR dy, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelAsinhGrad<DTYPE_Y> op;
    op.Init(y, dy, z, tiling_data.inputNum, tiling_data.tileLength);
    op.Process();
}